package io_test

import (
	"bytes"
	"context"
	"testing"
	"time"

	"github.com/prometheus/prometheus/model/rulefmt"
	"github.com/stretchr/testify/assert"

	"github.com/slok/sloth/internal/log"
	"github.com/slok/sloth/internal/storage/io"
	"github.com/slok/sloth/pkg/common/model"
)

func TestGroupedRulesYAMLRepoStore(t *testing.T) {
	tests := map[string]struct {
		slos    model.PromSLOGroupResult
		expYAML string
		expErr  bool
	}{
		"Having 0 SLO rules should fail.": {
			slos:   model.PromSLOGroupResult{},
			expErr: true,
		},

		"Having 0 SLO rules generated should fail.": {
			slos: model.PromSLOGroupResult{
				SLOResults: []model.PromSLOResult{},
			},
			expErr: true,
		},

		"Having a single SLI recording rule should render correctly.": {
			slos: model.PromSLOGroupResult{SLOResults: []model.PromSLOResult{
				{
					SLO: model.PromSLO{ID: "test1"},
					PrometheusRules: model.PromSLORules{
						SLIErrorRecRules: model.PromRuleGroup{
							Name: "sloth-slo-sli-recordings-test1",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record",
									Expr:   "test-expr",
									Labels: map[string]string{"test-label": "one"},
								},
							}},
					},
				},
			}},
			expYAML: `
---
# Code generated by Sloth (dev): https://github.com/slok/sloth.
# DO NOT EDIT.

groups:
- name: sloth-slo-sli-recordings-test1
  rules:
  - record: test:record
    expr: test-expr
    labels:
      test-label: one
`,
		},
		"Having a single metadata recording rule should render correctly.": {
			slos: model.PromSLOGroupResult{SLOResults: []model.PromSLOResult{
				{
					SLO: model.PromSLO{ID: "test1"},
					PrometheusRules: model.PromSLORules{
						MetadataRecRules: model.PromRuleGroup{
							Name: "sloth-slo-meta-recordings-test1",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record",
									Expr:   "test-expr",
									Labels: map[string]string{"test-label": "one"},
								},
							}},
					},
				},
			}},
			expYAML: `
---
# Code generated by Sloth (dev): https://github.com/slok/sloth.
# DO NOT EDIT.

groups:
- name: sloth-slo-meta-recordings-test1
  rules:
  - record: test:record
    expr: test-expr
    labels:
      test-label: one
`,
		},
		"Having a single SLO alert rule should render correctly.": {
			slos: model.PromSLOGroupResult{SLOResults: []model.PromSLOResult{
				{
					SLO: model.PromSLO{ID: "test1"},
					PrometheusRules: model.PromSLORules{
						AlertRules: model.PromRuleGroup{
							Name:     "sloth-slo-alerts-test1",
							Interval: 42 * time.Minute,
							Rules: []rulefmt.Rule{
								{
									Alert:       "testAlert",
									Expr:        "test-expr",
									Labels:      map[string]string{"test-label": "one"},
									Annotations: map[string]string{"test-annot": "one"},
								},
							}},
					},
				},
			}},
			expYAML: `
---
# Code generated by Sloth (dev): https://github.com/slok/sloth.
# DO NOT EDIT.

groups:
- name: sloth-slo-alerts-test1
  interval: 42m
  rules:
  - alert: testAlert
    expr: test-expr
    labels:
      test-label: one
    annotations:
      test-annot: one
`,
		},

		"Having a mixed example of multiple rules and options should render correctly.": {
			slos: model.PromSLOGroupResult{SLOResults: []model.PromSLOResult{
				{
					SLO: model.PromSLO{ID: "testa"},
					PrometheusRules: model.PromSLORules{
						SLIErrorRecRules: model.PromRuleGroup{
							Name: "sloth-slo-sli-recordings-testa",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-a1",
									Expr:   "test-expr-a1",
									Labels: map[string]string{"test-label": "a-1"},
								},
								{
									Record: "test:record-a2",
									Expr:   "test-expr-a2",
									Labels: map[string]string{"test-label": "a-2"},
								},
							}},
						MetadataRecRules: model.PromRuleGroup{
							Name: "sloth-slo-meta-recordings-testa",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-a3",
									Expr:   "test-expr-a3",
									Labels: map[string]string{"test-label": "a-3"},
								},
								{
									Record: "test:record-a4",
									Expr:   "test-expr-a4",
									Labels: map[string]string{"test-label": "a-4"},
								},
							}},
						AlertRules: model.PromRuleGroup{
							Name:     "sloth-slo-alerts-testa",
							Interval: 15 * time.Minute, // Custom interval.
							Rules: []rulefmt.Rule{
								{
									Alert:       "testAlertA1",
									Expr:        "test-expr-a1",
									Labels:      map[string]string{"test-label": "a-1"},
									Annotations: map[string]string{"test-annot": "a-1"},
								},
								{
									Alert:       "testAlertA2",
									Expr:        "test-expr-a2",
									Labels:      map[string]string{"test-label": "a-2"},
									Annotations: map[string]string{"test-annot": "a-2"},
								},
							}},
					},
				},
				{
					SLO: model.PromSLO{ID: "testb"},
					PrometheusRules: model.PromSLORules{
						SLIErrorRecRules: model.PromRuleGroup{
							Name: "sloth-slo-sli-recordings-testb",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-b1",
									Expr:   "test-expr-b1",
									Labels: map[string]string{"test-label": "b-1"},
								},
							}},
						MetadataRecRules: model.PromRuleGroup{
							Name: "sloth-slo-meta-recordings-testb",
							Rules: []rulefmt.Rule{
								{
									Record: "test:record-b2",
									Expr:   "test-expr-b2",
									Labels: map[string]string{"test-label": "b-2"},
								},
							}},
						AlertRules: model.PromRuleGroup{
							Name: "sloth-slo-alerts-testb",
							Rules: []rulefmt.Rule{
								{
									Alert:       "testAlertB1",
									Expr:        "test-expr-b1",
									Labels:      map[string]string{"test-label": "b-1"},
									Annotations: map[string]string{"test-annot": "b-1"},
								},
							}},
						ExtraRules: []model.PromRuleGroup{
							{
								Name:     "sloth-slo-extra-rules-000-testb",
								Interval: 42 * time.Minute, Rules: []rulefmt.Rule{
									{
										Alert:       "testAlertZ1",
										Expr:        "test-expr-z1",
										Labels:      map[string]string{"test-label": "z-1"},
										Annotations: map[string]string{"test-annot": "z-1"},
									},
								}},
							{}, // Should be skipped.
							{
								Name: "sloth-slo-extra-rules-001-testb",
								Rules: []rulefmt.Rule{
									{
										Alert:       "testAlertZ2",
										Expr:        "test-expr-z2",
										Labels:      map[string]string{"test-label": "z-2"},
										Annotations: map[string]string{"test-annot": "z-2"},
									},
									{
										Alert:       "testAlertZ3",
										Expr:        "test-expr-z3",
										Labels:      map[string]string{"test-label": "z-3"},
										Annotations: map[string]string{"test-annot": "z-3"},
									},
								},
							},
						},
					},
				},
			}},
			expYAML: `
---
# Code generated by Sloth (dev): https://github.com/slok/sloth.
# DO NOT EDIT.

groups:
- name: sloth-slo-sli-recordings-testa
  rules:
  - record: test:record-a1
    expr: test-expr-a1
    labels:
      test-label: a-1
  - record: test:record-a2
    expr: test-expr-a2
    labels:
      test-label: a-2
- name: sloth-slo-meta-recordings-testa
  rules:
  - record: test:record-a3
    expr: test-expr-a3
    labels:
      test-label: a-3
  - record: test:record-a4
    expr: test-expr-a4
    labels:
      test-label: a-4
- name: sloth-slo-alerts-testa
  interval: 15m
  rules:
  - alert: testAlertA1
    expr: test-expr-a1
    labels:
      test-label: a-1
    annotations:
      test-annot: a-1
  - alert: testAlertA2
    expr: test-expr-a2
    labels:
      test-label: a-2
    annotations:
      test-annot: a-2
- name: sloth-slo-sli-recordings-testb
  rules:
  - record: test:record-b1
    expr: test-expr-b1
    labels:
      test-label: b-1
- name: sloth-slo-meta-recordings-testb
  rules:
  - record: test:record-b2
    expr: test-expr-b2
    labels:
      test-label: b-2
- name: sloth-slo-alerts-testb
  rules:
  - alert: testAlertB1
    expr: test-expr-b1
    labels:
      test-label: b-1
    annotations:
      test-annot: b-1
- name: sloth-slo-extra-rules-000-testb
  interval: 42m
  rules:
  - alert: testAlertZ1
    expr: test-expr-z1
    labels:
      test-label: z-1
    annotations:
      test-annot: z-1
- name: sloth-slo-extra-rules-001-testb
  rules:
  - alert: testAlertZ2
    expr: test-expr-z2
    labels:
      test-label: z-2
    annotations:
      test-annot: z-2
  - alert: testAlertZ3
    expr: test-expr-z3
    labels:
      test-label: z-3
    annotations:
      test-annot: z-3
`,
		},
	}

	for name, test := range tests {
		t.Run(name, func(t *testing.T) {
			assert := assert.New(t)

			var gotYAML bytes.Buffer
			repo := io.NewStdPrometheusGroupedRulesYAMLRepo(&gotYAML, log.Noop)
			err := repo.StoreSLOs(context.TODO(), test.slos)

			if test.expErr {
				assert.Error(err)
			} else if assert.NoError(err) {
				assert.Equal(test.expYAML, gotYAML.String())
			}
		})
	}
}
